# -*- coding: utf-8 -*-
import torch
import time
import pandas as pd
import json
import numpy as np
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
    EvalPrediction,
)

class ERNIEFineTuner:
    def __init__(self, model, tokenizer, train_dataset, eval_dataset, compute_metrics, training_args):
        self.trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            tokenizer=tokenizer,
            compute_metrics=compute_metrics,
        )

    def train(self):
        print("微调开始...")
        train_result = self.trainer.train()
        print("微调完成！")
        return train_result

    def evaluate(self):
        print("评估开始...")
        eval_metrics = self.trainer.evaluate()
        print("评估完成！")
        return eval_metrics